# models/model_manager.py
import torch
from typing import Dict, List
from pathlib import Path
import shutil

class ModelManager:
    """模型管理器"""
    
    def __init__(self, models_dir: str = "data/models"):
        self.models_dir = Path(models_dir)
        self.models_dir.mkdir(parents=True, exist_ok=True)
        
        self.model_registry = self._load_registry()
    
    def _load_registry(self) -> Dict:
        """加载模型注册表"""
        registry_file = self.models_dir / "model_registry.json"
        if registry_file.exists():
            with open(registry_file, 'r') as f:
                return json.load(f)
        return {}
    
    def _save_registry(self):
        """保存模型注册表"""
        registry_file = self.models_dir / "model_registry.json"
        with open(registry_file, 'w') as f:
            json.dump(self.model_registry, f, indent=2)
    
    def register_model(self, model_name: str, model_path: str, metadata: Dict):
        """注册模型"""
        self.model_registry[model_name] = {
            "path": model_path,
            "metadata": metadata,
            "registered_at": datetime.now().isoformat()
        }
        self._save_registry()
    
    def load_model(self, model_name: str, model_class=None):
        """加载模型"""
        if model_name not in self.model_registry:
            raise ValueError(f"模型 {model_name} 未注册")
        
        model_info = self.model_registry[model_name]
        model_path = self.models_dir / model_info["path"]
        
        if model_class:
            model = model_class()
            model.load_state_dict(torch.load(model_path))
            return model
        else:
            return torch.load(model_path)
    
    def get_best_model(self, metric: str = "performance") -> str:
        """获取最佳模型"""
        best_model = None
        best_score = -float('inf')
        
        for model_name, info in self.model_registry.items():
            score = info["metadata"].get(metric, 0)
            if score > best_score:
                best_score = score
                best_model = model_name
        
        return best_model
    
    def cleanup_models(self, keep_top_k: int = 5):
        """清理模型，只保留前K个"""
        # 按性能排序
        sorted_models = sorted(
            self.model_registry.items(),
            key=lambda x: x[1]["metadata"].get("performance", 0),
            reverse=True
        )
        
        # 删除排名靠后的模型
        for i, (model_name, info) in enumerate(sorted_models):
            if i >= keep_top_k:
                model_path = self.models_dir / info["path"]
                if model_path.exists():
                    model_path.unlink()
                del self.model_registry[model_name]
        
        self._save_registry()
    
    def export_model(self, model_name: str, export_path: str):
        """导出模型"""
        if model_name not in self.model_registry:
            raise ValueError(f"模型 {model_name} 未注册")
        
        model_info = self.model_registry[model_name]
        source_path = self.models_dir / model_info["path"]
        
        # 复制模型文件
        shutil.copy2(source_path, export_path)
        
        # 复制元数据
        metadata_path = Path(export_path).with_suffix('.json')
        with open(metadata_path, 'w') as f:
            json.dump(model_info["metadata"], f, indent=2)